# LangGraph 多工具智能体 - 改进版（使用 LLM 原生工具调用）
# 无需手动意图识别！LLM 自动判断并调用工具

from typing import Annotated, Literal
from typing_extensions import TypedDict
from langgraph.graph import StateGraph, START, END
from langgraph.graph.message import add_messages
from langgraph.prebuilt import ToolNode
from langchain.chat_models import init_chat_model
from langchain.schema import AIMessage
from langchain.tools import tool
import re


# ========== 1. 定义工具（与之前相同）==========
@tool
def get_weather(city: str) -> str:
    """
    查询指定城市的天气信息
    
    Args:
        city: 城市名称，如：北京、上海、深圳
    
    Returns:
        该城市的天气信息
    """
    weather_data = {
        "北京": "晴天，气温 18-28℃，东南风 2-3级，空气质量良",
        "上海": "多云转晴，气温 22-30℃，东风 3-4级，空气质量优",
        "深圳": "雷阵雨，气温 26-32℃，南风 2级，空气质量良",
        "广州": "阴转小雨，气温 25-31℃，无持续风向，空气质量良",
        "杭州": "晴，气温 20-29℃，东北风 1-2级，空气质量优"
    }
    result = weather_data.get(city, f"暂无{city}的天气数据，仅支持：北京、上海、深圳、广州、杭州")
    return f"{city}的天气：{result}"


@tool
def calculator(expression: str) -> str:
    """
    计算数学表达式
    
    Args:
        expression: 数学表达式，如：2+3、10*5、100/4、(2+3)*5
    
    Returns:
        计算结果
    """
    try:
        clean_expr = expression.replace(" ", "")
        if not re.match(r'^[\d+\-*/().]+$', clean_expr):
            return "错误：表达式包含非法字符，仅支持数字和 +-*/() 运算符"
        
        result = eval(clean_expr)
        return f"{expression} = {result}"
    except Exception as e:
        return f"计算错误：{str(e)}"


# ========== 2. 定义 State（更简洁！）==========
class AgentState(TypedDict):
    """智能体状态 - 只需要消息历史！"""
    messages: Annotated[list, add_messages]


# ========== 3. 初始化 LLM 并绑定工具 ==========
llm = init_chat_model(model="deepseek-chat", model_provider="deepseek")

# 关键：将工具绑定到 LLM
tools = [get_weather, calculator]
llm_with_tools = llm.bind_tools(tools)

print("✅ LLM 已绑定工具：", [tool.name for tool in tools])


# ========== 4. 定义节点（只需要2个！）==========

# 节点1：调用 LLM（带工具）
def agent_node(state: AgentState) -> dict:
    """
    调用 LLM，LLM 会自动判断是否需要调用工具
    
    返回：
    - 如果需要工具：返回带 tool_calls 的 AIMessage
    - 如果不需要：返回普通的 AIMessage（直接回答）
    """
    messages = state["messages"]
    response = llm_with_tools.invoke(messages)
    
    # 打印调试信息
    if hasattr(response, 'tool_calls') and response.tool_calls:
        print(f"[Agent] LLM 决定调用工具: {[tc['name'] for tc in response.tool_calls]}")
    else:
        print(f"[Agent] LLM 直接回答，无需工具")
    
    return {"messages": [response]}


# 节点2：工具执行节点（LangGraph 内置！）
# ToolNode 会自动执行 AIMessage 中的 tool_calls
tool_node = ToolNode(tools)


# ========== 5. 定义路由函数（检查是否需要调用工具）==========
def should_continue(state: AgentState) -> Literal["tools", "end"]:
    """
    检查最后一条消息是否包含工具调用
    
    Returns:
        "tools": 需要执行工具
        "end": 无需工具，直接结束
    """
    last_message = state["messages"][-1]
    
    # 检查是否有工具调用
    if hasattr(last_message, 'tool_calls') and last_message.tool_calls:
        return "tools"
    
    return "end"


# ========== 6. 构建图（超简洁！）==========
graph_builder = StateGraph(AgentState)

# 添加节点
graph_builder.add_node("agent", agent_node)      # LLM 节点
graph_builder.add_node("tools", tool_node)       # 工具节点

# 定义边
graph_builder.add_edge(START, "agent")           # 开始 → Agent

# 条件边：Agent 之后检查是否需要工具
graph_builder.add_conditional_edges(
    "agent",
    should_continue,
    {
        "tools": "tools",  # 需要工具 → 工具节点
        "end": END         # 不需要 → 结束
    }
)

# 工具执行完后，返回 Agent（让 LLM 基于工具结果生成回答）
graph_builder.add_edge("tools", "agent")

# 编译图
improved_graph = graph_builder.compile()


# ========== 7. 可视化（可选）==========
try:
    png_data = improved_graph.get_graph().draw_mermaid_png()
    with open("/Users/baimu/PycharmProjects/2505A/boss-llm/boss/utils/langgraph/graph-improved.png", "wb") as f:
        f.write(png_data)
    print("✅ 流程图已保存: graph-improved.png\n")
except Exception as e:
    print(f"⚠️  可视化失败: {e}\n")


# ========== 8. 运行示例 ==========
def run_improved_agent():
    """运行改进版智能体"""
    print("=" * 60)
    print("🤖 改进版多工具智能体（使用 LLM 原生工具调用）")
    print("=" * 60)
    print("功能：")
    print("  📍 天气查询：自动识别城市并查询天气")
    print("  🔢 计算器：自动识别数学表达式并计算")
    print("  💬 普通对话：直接回答，无需工具")
    print("\n提示：输入空字符串或 'quit' 退出\n")
    
    while True:
        user_input = input("👤 User: ")
        
        if user_input.strip() == "" or user_input.lower() == "quit":
            print("\n👋 对话结束！")
            break
        
        # 初始化状态
        initial_state = {
            "messages": [{"role": "user", "content": user_input}]
        }
        
        # 执行图
        print("🤖 Assistant: ", end="", flush=True)
        
        for event in improved_graph.stream(initial_state):
            for node_name, node_output in event.items():
                if node_name == "agent" and "messages" in node_output:
                    latest_msg = node_output["messages"][-1]
                    if isinstance(latest_msg, AIMessage) and latest_msg.content:
                        # 只打印最终的文本回答
                        if not hasattr(latest_msg, 'tool_calls') or not latest_msg.tool_calls:
                            print(latest_msg.content)
        
        print()  # 换行


# ========== 9. 启动 ==========
if __name__ == "__main__":
    run_improved_agent()

